package matchbox;

import java.text.Format;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.Random;

import javax.sql.rowset.Predicate;

import probability.GaussianDistribution;
import probability.ProbabilityDistribution;
import utils.Pair;
import utils.Printer;
import inference.Graph;
import inference.GraphEdge;
import inference.GraphFactor;
import inference.MessagePassing;
import inference.GraphFactor.Type;
import inference.GraphVariable;
import basics.DenseVector;
import basics.FeaturesSet;
import basics.Vector;
import basics.VectorMatrix;
import basics.VectorMatrixUtils;
import basics.VectorMatrixUtils.OperationOpt;

/**
 * @author Ehsan
 * 
 */
public class Matchbox extends Classifier {

	private int _k = 1;
	private Graph _factorGraph;
	private GaussianDistribution _pu;
	private GaussianDistribution _pv;

	public static Graph buildGraph(int lDimension) {
		GraphEdge fzEdge = new GraphEdge();
		Graph fg = new Graph();
		GraphFactor fzbr = new ZBRFactor(NodeNames.F_ZBR);
		GraphVariable r = new GraphVariable(NodeNames.R);
		RRFactor frtr = new RRFactor(NodeNames.F_RtR);

		for (int k = 0; k < lDimension; k++) {
			STFactor fs = new STFactor(NodeNames.F_S + k);
			STFactor ft = new STFactor(NodeNames.F_T + k);
			GraphVariable s = new GraphVariable(NodeNames.S + k);
			GraphVariable t = new GraphVariable(NodeNames.T + k);
			fs.edges().put(fs.name, new GraphEdge(s));
			ft.edges().put(ft.name, new GraphEdge(t));
			s.edges().put(fs.name, new GraphEdge(fs));
			t.edges().put(ft.name, new GraphEdge(ft));

			frtr.addST(fs, ft);

			GraphVariable z = new GraphVariable(NodeNames.Z + k);

			GraphFactor fz = new ZFactor(NodeNames.F_Z + k);
			fz.edges().put(s.name, new GraphEdge(s));
			fz.edges().put(t.name, new GraphEdge(t));
			fz.edges().put(z.name, new GraphEdge(z));

			z.edges().put(NodeNames.F_Z, new GraphEdge(fz));
			z.edges().put(NodeNames.F_ZBR, new GraphEdge(fzbr));

			fzEdge.nodes().add(z);
		}

		fzbr.edges().put(NodeNames.Z, fzEdge);

		// STFactor fb = new STFactor(NodeNames.F_B);
		// GraphVariable b = new GraphVariable(NodeNames.B);
		// fb.edges().put(fb.name, new GraphEdge(b));
		// b.edges().put(fb.name, new GraphEdge(fb));
		// fzbr.edges().put(b.name, new GraphEdge(b));

		GraphVariable rt = new GraphVariable(NodeNames.RT);
		rt.edges().put(NodeNames.F_ZBR, new GraphEdge(fzbr));

		frtr.edges().put(NodeNames.RT, new GraphEdge(rt));
		fzbr.edges().put(NodeNames.RT, new GraphEdge(rt)); // frtr
		rt.edges().put(NodeNames.F_RtR, new GraphEdge(frtr));

		r.edges().put(NodeNames.F_RtR, new GraphEdge(frtr));
		frtr.edges().put(NodeNames.R, new GraphEdge(r));

		fg.queryNodes().add(r);
		// fg.evidenceNodes().add(ft);
		// fg.evidenceNodes().add(fs);

		return fg;
	}

	private static GraphFactor GeneratePreprocessedUserItemFactor(String name, FeaturesSet X, int lDimension) {
		STFactor f = new STFactor(name);
		// f.Preprocess(X, lDimension);
		return f;
	}

	private static void preprocess(VectorMatrix xy, int lDim, Graph g, String name) {
		for (int k = 0; k < lDim; k++) {
			STFactor fs = (STFactor) g.node(name + k);
			fs.Preprocess(xy);
		}
	}

	private static void postprocess(int lDim, Graph g) {
		for (int k = 0; k < lDim; k++) {
			ZFactor fz = ((ZFactor) g.node(NodeNames.F_Z + k));
			((STFactor) g.node(NodeNames.F_S + k)).updateLatent((GaussianDistribution) fz.toCached(
					g.node(NodeNames.S + k)).getDistribution());
			((STFactor) g.node(NodeNames.F_T + k)).updateLatent((GaussianDistribution) fz.toCached(
					g.node(NodeNames.T + k)).getDistribution());
		}
		// ((STFactor) g.node(NodeNames.F_B)).updateLatent((GaussianDistribution) ((ZBRFactor) g.node(NodeNames.F_ZBR))
		// .toCached(g.node(NodeNames.B)).getDistribution());
	}

	public void train(FeaturesSet X, FeaturesSet Y, FeaturesSet R, int lDimension) {
		_factorGraph = buildGraph(lDimension);
		_factorGraph.collectNodes();

		final int count = 5;

		for (int iteration = 1; iteration <= count; iteration++) {
			train(X, Y, R, lDimension, _factorGraph);
		}
	}

	private void train(FeaturesSet X, FeaturesSet Y, FeaturesSet R, int lDimension, Graph g) {
		_k = lDimension;

		((RRFactor) g.node(NodeNames.F_RtR)).setBounds(R.min(VectorMatrixUtils.OperationOpt.NonZero),
				R.max(VectorMatrixUtils.OperationOpt.NonZero));

		// for (int iteration = 0; iteration < 5; iteration++) {
		for (int i = 0; i < X.size(); i++) {
			ProbabilityDistribution pr = null;

			// preprocess variables s, t
			DenseVector x = (DenseVector) X.get(i);

			// preprocess(x, lDimension, g, NodeNames.F_S);
			for (int j = 0; j < Y.size(); j++) {
				if (R.get(i).get(j) <= 0) {
					continue;
				}
				DenseVector y = (DenseVector) Y.get(j);
				// y = (DenseVector) y.concat(new DenseVector(new double[] { 1. }));
				// preprocess(y, lDimension, g, NodeNames.F_T);
				// PREPROCESS IS SUPPOSED TO BE DONE IN trainUserItem. OTHERWISE, TAKE CARE OF THAT!
				Printer.print("------------- Start Iteration <" + i + "," + j + ">--------------");
				trainUserItem(x, y, R.get(i).get(j));
				Printer.print("------------- End Iteration <" + i + "," + j + ">--------------");
				postprocess(lDimension, g);
			}
		}
	}

	private Vector _prev_u = null, _prev_it = null;

	protected void trainUserItem(VectorMatrix u, VectorMatrix it, double r) {
		boolean process_b = false;
		if (_prev_u != (Vector) u) {
			_prev_u = (Vector) u;
			preprocess(u, _k, _factorGraph, NodeNames.F_S);
			process_b = true;
		}
		if (_prev_it != (Vector) it) {
			_prev_it = (Vector) it;
			preprocess(it, _k, _factorGraph, NodeNames.F_T);
			process_b = true;
		}

		// if (process_b) {
		// Vector uv = _prev_u.concat(_prev_it);
		// ((STFactor) _factorGraph.node(NodeNames.F_B)).Preprocess(uv);
		// }

		Graph g = _factorGraph;
		((RRFactor) g.node(NodeNames.F_RtR)).setObservedRate(r);

		pass(r, _k, g, null);
		// pass(r, _k, g, null, true);
		// epPass();
	}

	private void epPass() {
		final int count = 0;
		Graph g = _factorGraph;
		for (int iteration = 1; iteration <= count; iteration++) {
			epZST(NodeNames.Z);
			((ZBRFactor) g.node(NodeNames.F_ZBR)).ep(g.node(NodeNames.RT));
			((RRFactor) g.node(NodeNames.F_RtR)).ep(g.node(NodeNames.R));
			epZST(NodeNames.S);
			epZST(NodeNames.T);
			postprocess(_k, g);
		}
	}

	private void epZST(String name) {
		for (int k = 0; k < _k; k++) {
			((ZFactor) _factorGraph.node(NodeNames.F_Z + k)).ep(_factorGraph.node(name + k));
			// if (name.equals(NodeNames.S)) {
			// ((STFactor) _factorGraph.node(NodeNames.F_Z)).ep(null); // exceptionally STFactor works with null
			// } else if (name.equals(NodeNames.T)) {
			// ((STFactor) _factorGraph.node(NodeNames.F_T)).ep(null); // exceptionally STFactor works with null
			// }
		}
	}

	private void pass(double r, int lDimension, Graph g, String var) {
		ProbabilityDistribution pr = null, pr_prev = null;
		int n = 0;
		do {
			n++;
			pr_prev = pr != null ? pr.clone() : null;
			pr = ((GraphVariable) g.node(NodeNames.R)).to(null, var).getDistribution();
			if (!((pr_prev == null || pr.getParam(0).minus(pr_prev.getParam(0)).sumAllElements() > .01) && n < 200))
				break;

			Printer.print("Forward finished >>> :\n" + pr.toString() + " real R: " + r);
			String other = var != null && var.equals(NodeNames.S) ? NodeNames.T : NodeNames.S;

			// backward ...
			for (int k = 0; k < lDimension; k++) {
				ZFactor fz = ((ZFactor) g.node(NodeNames.F_Z + k));
				fz.to(((GraphVariable) g.node(other + k)), var); // NodeNames.S
			}
			Printer.print("Backward finished\n");
			epPass();
		} while (true);
		Printer.print("While finished\n");
	}

	private void cachePuPv() {
		_pu = null;
		_pv = null;
		for (int k = 0; k < _k; k++) {
			GaussianDistribution t = ((STFactor) _factorGraph.node(NodeNames.F_S + k)).getPuv();
			if (_pu == null) {
				_pu = t;
			} else {
				GaussianDistribution.product(_pu, t);
			}
			t = ((STFactor) _factorGraph.node(NodeNames.F_T + k)).getPuv();
			if (_pv == null) {
				_pv = t;
			} else {
				GaussianDistribution.product(_pv, t);
			}
		}
	}

	public double infer(VectorMatrix user, VectorMatrix item) {
		// int result = 1;

		// user = ((DenseVector) user).concat(new DenseVector(new double[] { 1. }));
		// item = ((DenseVector) item).concat(new DenseVector(new double[] { 1. }));

		preprocess(user, _k, _factorGraph, NodeNames.F_S);
		preprocess(item, _k, _factorGraph, NodeNames.F_T);

		GaussianDistribution pr = (GaussianDistribution) ((GraphVariable) _factorGraph.node(NodeNames.R)).to(null)
				.getDistribution();
		// pr = GaussianDistribution.product(pr, _pu, _pv);
		// for (int i = 2; i < 10; i++) {
		// if (pr.measure(i) > pr.measure(result)) {
		// result = i;
		// }
		// }

		return pr.getMu(); // result;
	}

	// public static String trainEval(FeaturesSet X, FeaturesSet Y, FeaturesSet R, FeaturesSet X_eval, FeaturesSet
	// Y_eval,
	// FeaturesSet R_eval, final int k) {
	//
	// double mean_r = 0d;
	// for (int i = 0; i < X.size(); i++) {
	// mean_r += VectorMatrixUtils.mean(R.get(i), OperationOpt.NonZero);
	// }
	// mean_r /= (double) X.size();
	//
	// HashMap<Integer, Double> mean_ir = new HashMap<Integer, Double>();
	// for (int j = 0; j < Y.size(); j++) {
	// double s = 0;
	// double n = 0;
	// for (int i = 0; i < X.size(); i++) {
	// if (R.get(i).get(j) > 0) {
	// s += R.get(i).get(j);
	// n++;
	// }
	// }
	// if (n > 0) {
	// mean_ir.put(j, s / n);
	// } else {
	// mean_ir.put(j, 0D);
	// }
	// }
	//
	// Matchbox m = new Matchbox();
	// m.train(X, Y, R, k);
	//
	// int correct = 0;
	// Double mae = 0.;
	// double mae_u_mean = 0;
	// double mae_i_mean = 0;
	// Double rmse = 0.;
	// double rmse_u_mean = 0;
	// double rmse_i_mean = 0;
	// int n = 0;
	// for (int i = 0; i < X_eval.size(); i++) {
	// VectorMatrix u = X_eval.get(i);
	// // Pair<Double, Double> mae_rmse =
	// Pair<Integer, Pair<Double, Double>> mae_rmse = m.testUser(u, i, Y_eval, R_eval, 75. / 100.);
	// n += mae_rmse.getKey();
	// mae += mae_rmse.getValue().getKey();
	// rmse += mae_rmse.getValue().getValue();
	//
	// // mae += mae_rmse.getKey();
	// // rmse += mae_rmse.getValue();
	// // (err * err);
	// // for (int j = 0; j < Y_eval.size(); j++) {
	// // if (R_eval.get(i).get(j) <= 0)
	// // continue;
	// // VectorMatrix it = Y_eval.get(j);
	// // double res = m.infer(u, it);
	// // double err = Math.abs(res - R_eval.get(i).get(j));
	// // double mean_u_err = Math.abs(R_eval.get(i).get(j) - mean_r);
	// // double mean_i_err = Math.abs(R_eval.get(i).get(j) - mean_ir.get(j));
	// // System.out.println("result: " + res + ", trueval " + R_eval.get(i).get(j) + ", err " + err
	// // + ", Mean_u_err " + mean_u_err + ", Mean_i_err " + mean_i_err);
	// //
	// // mae_u_mean += mean_u_err;
	// // mae_i_mean += mean_i_err;
	// // rmse_u_mean += (mean_u_err * mean_u_err);
	// // rmse_i_mean += (mean_i_err * mean_i_err);
	// // n++;
	// // if (err < 1.) {
	// // correct++;
	// // }
	// // }
	// }
	//
	// String str = String.format("Accuracy: %2.0f, MAE  = %f, Mean-u-MAE  = %f, Mean-i-MAE  = %f\n", (double) correct
	// * 100 / (double) n, mae / (double) n, mae_u_mean / (double) n, mae_i_mean / (double) n);
	// str = str
	// + String.format("              RMSE = %f, Mean-u-RMSE = %f, Mean-i-RMSE = %f", rmse / (double) n,
	// rmse_u_mean / (double) n, rmse_i_mean / (double) n);
	// return str;
	// }

	//
	// protected Pair<Integer, Pair<Double, Double>> testUser(VectorMatrix u, int userIdx, FeaturesSet y_eval,
	// FeaturesSet r_eval, double percent) {
	// Pair<Integer, Pair<Double, Double>> p = new Pair<Integer, Pair<Double, Double>>();
	// p.setKey(0);
	// p.setValue(new Pair<Double, Double>(0., 0.));
	//
	// int rated = 0;
	//
	// for (int j = 0; j < y_eval.size(); j++) {
	// if (r_eval.get(userIdx).get(j) > 0) {
	// rated++;
	// }
	// }
	//
	// if (rated < 1)
	// return p;
	//
	// preprocess(u, _k, _factorGraph, NodeNames.F_S);
	// int test = 0;
	// for (int j = 0, n = 0; j < y_eval.size(); j++) {
	// if (r_eval.get(userIdx).get(j) > 0) {
	// if (n < percent * rated) {
	// preprocess(y_eval.get(j), _k, _factorGraph, NodeNames.F_T);
	// trainUserItem(r_eval.get(userIdx).get(j));
	// } else {
	// double res = infer(u, y_eval.get(j));
	// double err = Math.abs(res - r_eval.get(userIdx).get(j));
	// // mae += err;
	// // rmse += (err * err);
	// test++;
	// p.setKey(test);
	// p.getValue().setKey(p.getValue().getKey() + err);
	// p.getValue().setValue(p.getValue().getValue() + (err * err));
	// }
	// n++;
	// }
	// }
	//
	// // p.setKey(p.getKey() / (rated - (int) percent * rated));
	// // p.setValue(p.getValue() / (rated - (int) percent * rated));
	//
	// return p;
	// }

	// public static void trainEval(FeaturesSet X, FeaturesSet Y, FeaturesSet R, final int k) {
	// final int trainig_size = 150;
	// final int test_size = 100;
	// FeaturesSet users = new FeaturesSet();
	// FeaturesSet ratings = new FeaturesSet();
	// FeaturesSet items = Y; // new FeaturesSet();
	// double mean_r = 0;
	//
	// System.out.println("Initializing dataset");
	// for (int i = 0; i < trainig_size; i++) {
	// if (i < X.size()) {
	// users.add(X.get(i));
	// }
	// // if (i < Y.size()) {
	// // items.add(Y.get(i));
	// // }
	// ratings.add(R.get(i));
	// mean_r += VectorMatrixUtils.mean(R.get(i), OperationOpt.NonZero);
	// }
	// mean_r = mean_r / (double) ratings.size();
	// System.out.println("Dataset initialized. Training with " + k + " latent dimensions ...");
	// Matchbox m = new Matchbox();
	// m.train(users, items, ratings, k);
	//
	// System.out.println("Training done! Evaluation ...");
	// int correct = 0;
	// double mae = 0;
	// int n = 0;
	// for (int i = 0; i < test_size; i++) {
	// VectorMatrix u = X.get(i + trainig_size);
	// for (int j = 0; j < items.size(); j++) {
	// if (R.get(i + trainig_size).get(j) <= 0)
	// continue;
	// VectorMatrix it = items.get(j);
	// double res = m.infer(u, it);
	// // res = 5. / (1 + Math.exp(-res));
	// double err = Math.abs(res - R.get(i + trainig_size).get(j));
	// System.out.println("result: " + res + ", trueval " + R.get(i + trainig_size).get(j) + ", err " + err);
	// mae += err;
	// n++;
	// if (err < 1.) {
	// correct++;
	// // System.out.println("correct!!!! " + correct + " till now!!!");
	// }
	// }
	// }
	//
	// System.out.println("Done!\nAccuracy = " + (double) correct / (double) n + " and MAE = " + mae / (double) n
	// + " while mean_r=" + mean_r);
	// }
}
